/*
 * Copyright 2012 Claude Houle claude.houle@gmail.com
 *
 * Licensed under the Apache License, Version 2.0 (the "License");
 * you may not use this file except in compliance with the License.
 * You may obtain a copy of the License at
 *
 * http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 */
package com.google.code.mochaccino.framework.crypto.keystore;

import com.google.code.mochaccino.framework.crypto.EncryptionException;
import com.google.common.base.Preconditions;
import java.io.BufferedInputStream;
import java.io.BufferedOutputStream;
import java.io.File;
import java.io.FileInputStream;
import java.io.FileNotFoundException;
import java.io.FileOutputStream;
import java.io.IOException;
import java.security.Key;
import java.security.KeyStore;
import java.security.KeyStoreException;
import java.security.NoSuchAlgorithmException;
import java.security.NoSuchProviderException;
import java.security.UnrecoverableKeyException;
import java.security.cert.Certificate;
import java.security.cert.CertificateException;
import java.util.ArrayList;
import java.util.Enumeration;
import java.util.List;
import org.apache.commons.io.IOUtils;

/**
 * The KeyStore Manager is a simple facade to easily manipulate keystores. Be sure
 * to call the load() method before any manipulation and save() after each manipulation.
 */
public final class KeyStoreManager {
	private File         file     = null;
	private KeyStoreType type     = null;
	private KeyStore     keyStore = null;
	private String       password = null;

	public KeyStoreManager( File file, KeyStoreType type, String password ) {
		this.file = Preconditions.checkNotNull( file );
		this.type = Preconditions.checkNotNull( type );
		this.password = Preconditions.checkNotNull( password );
	}

	/** Contain Alias? */
	public boolean containsAlias( String alias ) throws EncryptionException {
		try {
			return keyStore.containsAlias( alias );
		} catch ( KeyStoreException e ) {
			throw new EncryptionException( e );
		}
	}

	/** Delete Entry */
	public void deleteEntry( String alias ) throws EncryptionException {
		try {
			keyStore.deleteEntry( alias );
		} catch ( KeyStoreException e ) {
			throw new EncryptionException( e );
		}
	}

	/** Get Certificate */
	public Certificate getCertificate( String alias ) throws EncryptionException {
		try {
			return keyStore.getCertificate( alias );
		} catch ( KeyStoreException e ) {
			throw new EncryptionException( e );
		}
	}

	/** Return All Certificate Aliases */
	public String[] getCertificateAliases() throws EncryptionException {
		try {
			List<String> certificateAlias = new ArrayList<String>();
			Enumeration<String> aliases = keyStore.aliases();
			while ( aliases.hasMoreElements() ) {
				String alias = aliases.nextElement();
				if ( keyStore.isCertificateEntry( alias ) ) {
					certificateAlias.add( alias );
				}
			}
			return certificateAlias.toArray( new String[certificateAlias.size()] );
		} catch ( KeyStoreException e ) {
			throw new EncryptionException( e );
		}
	}

	/** Get Key */
	public Key getKey( String alias ) throws EncryptionException {
		return getKey( alias, password );
	}

	/** Get Key */
	public Key getKey( String alias, String keyPassword ) throws EncryptionException {
		try {
			return keyStore.getKey( alias, keyPassword.toCharArray() );
		} catch ( UnrecoverableKeyException e ) {
			throw new EncryptionException( e );
		} catch ( KeyStoreException e ) {
			throw new EncryptionException( e );
		} catch ( NoSuchAlgorithmException e ) {
			throw new EncryptionException( e );
		}
	}

	/** Get Key Aliases */
	public String[] getKeyAliases() throws EncryptionException {
		try {
			List<String> keyAlias = new ArrayList<String>();
			Enumeration<String> aliases = keyStore.aliases();
			while ( aliases.hasMoreElements() ) {
				String alias = aliases.nextElement();
				if ( keyStore.isKeyEntry( alias ) ) {
					keyAlias.add( alias );
				}
			}
			return keyAlias.toArray( new String[keyAlias.size()] );
		} catch ( KeyStoreException e ) {
			throw new EncryptionException( e );
		}
	}

	/** Get Certificate Chain */
	public Certificate[] getKeyCertificateChain( String alias ) throws EncryptionException {
		try {
			return keyStore.getCertificateChain( alias );
		} catch ( KeyStoreException e ) {
			throw new EncryptionException( e );
		}
	}

	public KeyStore getKeyStore() {
		return keyStore;
	}

	/** Return the password */
	public String getPassword() {
		return password;
	}

	public KeyStoreType getType() {
		return type;
	}

	public void load() throws EncryptionException {
		try {
			keyStore = type.createNewKeyStore();
			if ( file.exists() && file.canRead() ) {
				BufferedInputStream in = null;
				try {
					in = new BufferedInputStream( new FileInputStream( file ) );
					keyStore.load( in, password.toCharArray() );
				} finally {
					IOUtils.closeQuietly( in );
				}
			} else {
				keyStore.load( null );
			}
		} catch ( KeyStoreException e ) {
			throw new EncryptionException( e );
		} catch ( NoSuchProviderException e ) {
			throw new EncryptionException( e );
		} catch ( CertificateException e ) {
			throw new EncryptionException( e );
		} catch ( NoSuchAlgorithmException e ) {
			throw new EncryptionException( e );
		} catch ( FileNotFoundException e ) {
			throw new EncryptionException( e );
		} catch ( IOException e ) {
			throw new EncryptionException( e );
		}
	}

	public void renameTo( File newFile ) throws IOException {
		if ( !file.renameTo( newFile ) ) {
			throw new IOException( "KeyStore could not be renamed!" );
		}
	}

	/** Set Certificate */
	public void setCertificate( String alias, Certificate cert ) throws EncryptionException {
		try {
			keyStore.setCertificateEntry( alias, cert );
		} catch ( KeyStoreException e ) {
			throw new EncryptionException( e );
		}
	}

	/** Set Key */
	public void setKey( String alias, Key key, Certificate[] chain ) throws EncryptionException {
		setKey( alias, key, password, chain );
	}

	/** Set Key */
	public void setKey( String alias, Key key, String password, Certificate[] chain ) throws EncryptionException {
		try {
			keyStore.setKeyEntry( alias, key, password.toCharArray(), chain );
		} catch ( KeyStoreException e ) {
			throw new EncryptionException( e );
		}
	}

	public void setPassword( String newPassword ) throws EncryptionException {
		this.password = newPassword;
		this.save();
	}

	public void save() throws EncryptionException {
		try {
			BufferedOutputStream out = null;
			try {
				out = new BufferedOutputStream( new FileOutputStream( file ) );
				keyStore.store( out, password.toCharArray() );
			} finally {
				IOUtils.closeQuietly( out );
			}
		} catch ( KeyStoreException e ) {
			throw new EncryptionException( e );
		} catch ( IOException e ) {
			throw new EncryptionException( e );
		} catch ( NoSuchAlgorithmException e ) {
			throw new EncryptionException( e );
		} catch ( CertificateException e ) {
			throw new EncryptionException( e );
		}
	}
}
